#include "io/libAV.h"
#include "io/vsMedia.h"

#include <assert.h>
#include <QTimer>
#include <QFile>
#include <QFileInfo>
#include <QDir>

#include <math.h>

extern "C" {
    #include <libavformat/avformat.h>
    #include <libavutil/dict.h>
    #include <libavutil/pixfmt.h>
    #include <libswscale/swscale.h>
}

#if LIBAVCODEC_VERSION_MAJOR<55
#define VS_OLD_LAV 1
#endif

#ifdef VS_OLD_LAV
    #define AV_CODEC_ID_BMP CODEC_ID_BMP
    #define av_frame_alloc avcodec_alloc_frame
    #define av_frame_free  av_free
    #define AVPixelFormat PixelFormat
#endif


libAVctrl::libAVctrl(QObject *parent) : QObject(parent) { av_register_all(); }
libAVctrl::~libAVctrl() {}

bool libAVctrl::avError(int code, const QString &file)
{
    const int bufsize = 100;
    char buf[bufsize];
    av_strerror(code, buf, bufsize);
    vsLog::e(QObject::tr("libAV: \"%1\" @ %2").arg(QString(buf)).arg(file));

    return false;
}

bool libAVctrl::probeMedia(vsMediaFile &f)
{
    AVFormatContext *fmt_ctx = NULL;
    AVDictionaryEntry *tag = NULL;
    unsigned int i;

    int ret = avformat_open_input(&fmt_ctx, f.path.toUtf8().data(), NULL, NULL);

    if (ret < 0)
        return avError(ret, f.path);
    else
    {
        ret = avformat_find_stream_info(fmt_ctx, NULL);

        if (ret < 0)
        {
            avformat_close_input(&fmt_ctx);
            return avError(ret, f.path); // no point in continuing, can't get general tags too
        }
    }

    f.bitrate = fmt_ctx->bit_rate;

    if (fmt_ctx->duration != (int64_t)AV_NOPTS_VALUE)
        f.duration = (fmt_ctx->duration + 5000) / (AV_TIME_BASE / 1000); // magic numbers, magic numbers everywhere

    for (i = 0; i < fmt_ctx->nb_streams; i++)
    {
        AVStream *stream = fmt_ctx->streams[i];

        while ((tag = av_dict_get(stream->metadata, "", tag, AV_DICT_IGNORE_SUFFIX)))
            f.tags.insert(QString::fromUtf8(tag->key).toLower(), QString::fromUtf8(tag->value));

        AVCodec *codec = avcodec_find_decoder(stream->codec->codec_id);
        AVCodecContext *cctx = avcodec_alloc_context3(codec); // will return 0 if no codec

        if (cctx && cctx->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            if (stream->codec->extradata_size > 0)
            {
#ifdef VS_OLD_LAV
                cctx->extradata = (unsigned char*)av_malloc(stream->codec->extradata_size);
                memcpy(cctx->extradata, stream->codec->extradata, stream->codec->extradata_size);
#else
                cctx->extradata = (unsigned char*)av_memdup(stream->codec->extradata, stream->codec->extradata_size);
#endif
                cctx->extradata_size = stream->codec->extradata_size;
            }

            ret = avcodec_open2(cctx, codec, 0);

            if (ret) avError(ret, f.path);
            else
            {
                // found a video stream, trying to get a single frame for cover picture
                AVFrame *pic = av_frame_alloc();
                AVPacket pkt;
                memset(&pkt, 0, sizeof(pkt));
                av_init_packet(&pkt);
                int got_pic = 0;

                ret = av_seek_frame(fmt_ctx, i, stream->first_dts + (int64_t)(2.f * av_q2d(stream->time_base)), 0); // seeking by 2 seconds from the start
                ret = av_read_frame(fmt_ctx, &pkt);

                if (ret) avError(ret, f.path);
                else
                    ret = avcodec_decode_video2(cctx, pic, &got_pic, &pkt);

                if (ret < 0)
                    avError(ret, f.path);

                if (ret > 0 || got_pic)
                {
                    while (!got_pic) // that's kinda dangerous, but encoder may want more packets to make a full frame
                    {
                        av_read_frame(fmt_ctx, &pkt);
                        avcodec_decode_video2(cctx, pic, &got_pic, &pkt);
                    }

                    // now encode that frame (that may have any format) into a simple bitmap
                    AVCodec *encoder = avcodec_find_encoder(AV_CODEC_ID_BMP);
                    AVCodecContext *encCCtx = avcodec_alloc_context3(encoder);

                    encCCtx->codec      = encoder;
                    encCCtx->codec_id   = encoder->id;
                    encCCtx->codec_type = encoder->type;
                    encCCtx->pix_fmt    = encoder->pix_fmts[0];
                    encCCtx->width      = pic->width;
                    encCCtx->height     = pic->height;

                    ret = avcodec_open2(encCCtx, encoder, 0);

                    if (ret) avError(ret, f.path);
                    else
                    {
                        int got_packet = 0;
                        AVPacket imgData;
                        memset(&imgData, 0, sizeof(imgData));
                        av_init_packet(&imgData);

#ifdef VS_OLD_LAV
                        imgData.size = pic->width * pic->height *
                            av_get_bytes_per_sample((AVSampleFormat)encCCtx->pix_fmt) + 54 + 20; // 20 for safety
                        imgData.data = (unsigned char*)av_malloc(imgData.size);
                        memset(imgData.data, 0, imgData.size);
#endif

                        if (pic->format != encCCtx->pix_fmt) // which is true in most cases
                        {
                            // so we need to convert pixel format of frame before encoding it

                            struct SwsContext *img_convert_ctx = sws_getContext(pic->width, pic->height, (AVPixelFormat)pic->format,
                                pic->width, pic->height, encCCtx->pix_fmt, SWS_BICUBIC, NULL, NULL, NULL);

                            AVFrame* scaledPic = av_frame_alloc();
                            int scaledPicBufSize = avpicture_get_size(encCCtx->pix_fmt, pic->width, pic->height);
                            unsigned char *picture_buf = (unsigned char*)av_malloc(scaledPicBufSize);
                            avpicture_fill((AVPicture *)scaledPic, picture_buf, encCCtx->pix_fmt, pic->width, pic->height);
                            sws_scale(img_convert_ctx, (const uint8_t * const *)pic->data, pic->linesize, 0, pic->height,
                                      scaledPic->data, scaledPic->linesize);
#ifdef VS_OLD_LAV
                            ret = avcodec_encode_video(encCCtx, imgData.data, imgData.size, scaledPic);
                            got_pic = ret >= imgData.size - 20;
#else
                            ret = avcodec_encode_video2(encCCtx, &imgData, scaledPic, &got_packet);
#endif
                            av_frame_free(&scaledPic);
                            av_free(picture_buf);
                        }
                        else
                        {
#ifdef VS_OLD_LAV
                            ret = avcodec_encode_video(encCCtx, imgData.data, imgData.size, pic);
                            got_packet = ret >= imgData.size - 20;
#else
                            ret = avcodec_encode_video2(encCCtx, &imgData, pic, &got_packet);
#endif
                        }

                        if (ret)
                            avError(ret, f.path);
                        else
                            if (got_packet)
                            {
                                f.covers.append(QImage::fromData(imgData.data, imgData.size, "bmp"));
                                av_free_packet(&imgData);
                            }
                    }

                    avcodec_close(encCCtx);
                    av_free(encCCtx);
                }

                av_free_packet(&pkt);
            }

            avcodec_close(cctx);
            av_free(cctx);
            // TODO: maybe should free codec?
        }
    }

    // container tags have priority over tags of streams
    // qmap will overwrite values for tags with same name
    while ((tag = av_dict_get(fmt_ctx->metadata, "", tag, AV_DICT_IGNORE_SUFFIX)))
        f.tags.insert(QString::fromUtf8(tag->key).toLower(), QString::fromUtf8(tag->value));

    avformat_close_input(&fmt_ctx);
    return ret == 0;
}

void libAVctrl::getInfo(vsMediaFile *file)
{
    // TODO: multithreading if freezes ui
    assert(file != 0);

    file->tags.clear();
    file->covers.clear();
    file->duration = 0;
    file->audioBitrate = 0;
    file->videoBitrate = 0;
    file->cachedCover.clear();
    file->cachedCoverMime.clear();
    file->cachedCoverSize = QSize();

    probeMedia(*file);

    vsLog::s("Media file processed: " + file->path);

    if (file->covers.isEmpty())
    {
        // try to find a cover pic in folder
        QFileInfo fi(file->path);
        QDir fd(fi.absolutePath());

        QFileInfoList pics = fd.entryInfoList({"*.bmp", "*.png", "*.jpg", "*.jpeg", "*.tga"},
            QDir::Files | QDir::NoSymLinks | QDir::NoDotAndDotDot | QDir::Readable);

        if (!pics.isEmpty())
        {
            QList<QRegExp> res;

            res.append(QRegExp("folder"));
            res.append(QRegExp("cover"));

            if (!fi.baseName().isEmpty())
                res.append(QRegExp(fi.baseName()));

            if (!fd.dirName().isEmpty())
                res.append(QRegExp(fd.dirName()));

            if (file->tags.contains("album"))
                res.append(QRegExp(file->tags["album"]));

            if (!res.isEmpty())
                foreach (const QFileInfo &pic, pics)
                    for (int i = 0; i < res.size(); ++i)
                        if (pic.baseName().contains(res[i]))
                            file->covers.append(QImage(pic.absoluteFilePath()));
        }
    }

    file->tags.insert("mime", vsMedia::instance()->mime(file->ext));
    file->status = vsTypes::seeding; // should set this even on fail, won't be seeded otherwise
    file->onInfoUpdated();
}
